from sklearn.model_selection import StratifiedShuffleSplit

def _show_dist():
    _1 = train_df['ratio_bin'].value_counts(normalize=True).reset_index().rename(columns={"proportion":"Train"})
    _2 = val_df['ratio_bin'].value_counts(normalize=True).reset_index().rename(columns={"proportion":"Validation"})
    _3 = test_df['ratio_bin'].value_counts(normalize=True).reset_index().rename(columns={"proportion":"Test"})

    _dist = pd.merge(_1, _2, on="ratio_bin", how="left")
    _dist = pd.merge(_dist, _3, on="ratio_bin", how="left").sort_values(by="ratio_bin")
    _dist["ratio_bin"] = _dist["ratio_bin"]/10
    return _dist.rename(columns={"ratio_bin":"Apnea Samples Group"}).round(3)
    
splitter = StratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=42)

X = df_dist['file']
y = df_dist['ratio_bin']

# First split: train vs. temp (val + test)
for train_idx, temp_idx in splitter.split(X, y):
    train_df = df_dist.iloc[train_idx]
    temp_df = df_dist.iloc[temp_idx]
    
    # Second split: val vs. test (split temp again)
    splitter_2 = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42)

    # exclude ratio_bin group with just one element
    exclude_groups = [_bin for _bin, cond in (temp_df['ratio_bin'].value_counts()==1).items() if cond==True]
    temp_df = temp_df[temp_df['ratio_bin'].apply(lambda x: x not in exclude_groups)]
    
    X_temp = temp_df['file']
    y_temp = temp_df['ratio_bin']
    
    for val_idx, test_idx in splitter_2.split(X_temp, y_temp):
        val_df  = temp_df.iloc[val_idx]
        test_df = temp_df.iloc[test_idx]

_show_dist()

train_files = train_df['file'].values
valid_files = val_df['file'].values
test_files  = test_df['file'].values

def fast_filter_fn(min_std=0.006):
    """
    heuristic of pixel distribution, avoid empty Mel Spectrograms
    
    Do not hurt to set minimum 0.006
        - min std tracheal 0.006
        - min std mic      0.004
    """
    def inner_fn(image, label):
        std = tf.math.reduce_std(image)
        return std > min_std
    return inner_fn

#map_type = {0:'Normal', 1:'Hypopnea', 2:'MixedApnea', 3:'ObstructiveApnea', 4:'CentralApnea'}
def keep_only_0_or_3(mel_spec, label):
    return tf.logical_or(tf.equal(label[1], 0), tf.equal(label[1], 3))

BATCH_SIZE = 16 #32 #64#//2
GLOBAL_BATCH_SIZE = BATCH_SIZE*strategy.num_replicas_in_sync

############################################
input_data = "Tracheal" # "Tracheal", "Mic"
#input_data = "Mic"
if input_data == "Mic":
    _type = 0
else:
    _type = 1
############################################


# Load TFRecord Dataset
train_ds = (
    tf.data.Dataset.from_tensor_slices(train_files)
    .interleave(
        lambda x: tf.data.TFRecordDataset(x),
        cycle_length=tf.data.AUTOTUNE,       # Number of files to read in parallel
        num_parallel_calls=tf.data.AUTOTUNE, # Number of parallel calls for interleave
        deterministic=False                  # Allow non-deterministic order for speed
    )
    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    .filter(keep_only_0_or_3) # just obstructive apnea
    # combined structure t_images -> (mic, tracheal), t_labels -> (label, label_type) 
    .map(lambda t_images, t_labels: (t_images[_type], t_labels[0]), num_parallel_calls=tf.data.AUTOTUNE)
    .filter(fast_filter_fn(min_std=0.006))
    #.cache()
    .shuffle(5000)
    .batch(GLOBAL_BATCH_SIZE)
    .prefetch(tf.data.AUTOTUNE)
)

valid_ds = (
    tf.data.Dataset.from_tensor_slices(valid_files)
    .interleave(
        lambda x: tf.data.TFRecordDataset(x),
        cycle_length=tf.data.AUTOTUNE,       # Number of files to read in parallel
        num_parallel_calls=tf.data.AUTOTUNE, # Number of parallel calls for interleave
        deterministic=False                  # Allow non-deterministic order for speed
    )
    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    .filter(keep_only_0_or_3) # just obstructive apnea
    # combined structure t_images -> (mic, tracheal), t_labels -> (label, label_type) 
    .map(lambda t_images, t_labels: (t_images[_type], t_labels[0]), num_parallel_calls=tf.data.AUTOTUNE)
    .filter(fast_filter_fn(min_std=0.006))
    #.cache()
    .batch(GLOBAL_BATCH_SIZE)
    .prefetch(tf.data.AUTOTUNE)
)

test_ds = (
    tf.data.Dataset.from_tensor_slices(test_files)
    .interleave(
        lambda x: tf.data.TFRecordDataset(x),
        cycle_length=tf.data.AUTOTUNE,       # Number of files to read in parallel
        num_parallel_calls=tf.data.AUTOTUNE, # Number of parallel calls for interleave
        deterministic=False                  # Allow non-deterministic order for speed
    )
    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    .filter(keep_only_0_or_3) # just obstructive apnea
    # combined structure t_images -> (mic, tracheal), t_labels -> (label, label_type) 
    .map(lambda t_images, t_labels: (t_images[_type], t_labels[0]), num_parallel_calls=tf.data.AUTOTUNE)
    .filter(fast_filter_fn(min_std=0.006))
    #.cache() # to have a deterministic test_ds
    .batch(GLOBAL_BATCH_SIZE)
    .prefetch(tf.data.AUTOTUNE)
)

print(f"\nUsing {input_data}, batches {GLOBAL_BATCH_SIZE}\n")

t_images, t_labels = valid_ds.take(1).as_numpy_iterator().next()

f,axs = plt.subplots(nrows=4, ncols=4)

for i, ax in enumerate(axs.ravel()):

    ax.imshow(t_images[i][::-1], cmap="magma")
    ax.axis("off")

plt.show()

import math


def warmup_cosine_lr(base_lr, total_epochs, warmup_epochs=5, min_lr=0.0):
    """
    Compute the learning rate for a given epoch using warmup + cosine annealing.
    
    Args:
        epoch (int): Current epoch number (0-based).
        base_lr (float): The initial learning rate after warmup.
        total_epochs (int): Total number of training epochs.
        warmup_epochs (int): Number of warmup epochs at the start of training.
        min_lr (float): Minimum learning rate at the end of cosine annealing.
    
    Returns:
        float: Adjusted learning rate.
    """
    def inner_function(epoch): 
        if epoch < warmup_epochs:
            # Linear warmup
            return base_lr * (epoch + 1) / warmup_epochs
        else:
            # Cosine annealing
            cosine_epochs = total_epochs - warmup_epochs
            epoch_in_cosine = epoch - warmup_epochs
            cosine_decay = 0.5 * (1 + math.cos(math.pi * epoch_in_cosine / cosine_epochs))
            return min_lr + (base_lr - min_lr) * cosine_decay

    return inner_function

# Example

lr_scheduler = warmup_cosine_lr(base_lr=4e-4, total_epochs=35, warmup_epochs=8, min_lr=1e-6)

x = list(range(35))
y = [lr_scheduler(e) for e in x]

f,ax=plt.subplots(figsize=(8,4))
ax.plot(x,y)
ax.set_title("Cosine annealing Scheduler")
ax.set_xlabel("Epoch")
ax.set_ylabel("Learning Rate")

f,ax = plt.subplots(ncols=3, figsize=(12,4))

ax[0].set_title("Accuracy")
ax[1].set_title("F1 Score")
ax[2].set_title("Loss")

x_seg = list(range(1,len(history['accuracy'])+1))

ax[0].plot(x_seg, history['accuracy'], color="navy", label="accuracy")
ax[0].plot(x_seg, history['val_accuracy'], color="orange",  label="val_accuracy")

ax[1].plot(x_seg, history['f1_score'], color="navy",  label="f1-score")
ax[1].plot(x_seg, history['val_f1_score'], color="orange", label="val_f1-score")

ax[2].plot(x_seg, history['loss'], color="navy", label="loss")
ax[2].plot(x_seg, history['val_loss'], color="orange", label="val_loss")

for i in range(3):
    ax[i].legend()
    ax[i].set_xlabel("Epoch")
    ax[i].grid(True)
    #ax[i].set_xlim([1,11])

plt.tight_layout()
plt.show()

f.savefig("loss_trach_epoch.pdf", dpi=300, bbox_inches='tight')

# test_ds other can change if called multiple times, due to interleave, prefetch, a deterministic test_ds must use .cache()

y_true = []
y_pred_proba = []

for batch_x, batch_y in tqdm( test_ds ):
    
    _preds = model.predict(batch_x, verbose=0)
    y_pred_proba.extend(_preds.flatten())  # Binary threshold
    y_true.extend(batch_y.numpy().flatten())

import numpy as np
from sklearn.metrics import precision_recall_fscore_support

thresholds = np.linspace(0.0, 1.0, 101)
precision, recall, f1 = [], [], []

for t in thresholds:
    preds = (y_pred_proba >= t).astype(int)
    p, r, f, _ = precision_recall_fscore_support(y_true, preds, average='binary', zero_division=0)
    precision.append(p)
    recall.append(r)
    f1.append(f)

f,ax = plt.subplots(figsize=(6, 3))
ax.plot(thresholds, precision, label='Precision')
ax.plot(thresholds, recall, label='Recall')
ax.plot(thresholds, f1, label='F1 Score')

ax.axvline(0.3, color='navy', linestyle='--', label='threshold=0.3')

ax.axvline(0.5, color='gray', linestyle='--', label='Default threshold 0.5')
ax.set_xlabel('Threshold')
ax.set_ylabel('Score')
ax.set_title('Precision, Recall, and F1 Score vs Threshold')

plt.legend(loc='lower right')
plt.grid(lw=0.25)
plt.show()

#f.savefig("tr_metrics.pdf", dpi=300, bbox_inches='tight')

from sklearn.metrics import confusion_matrix

# True and predicted labels
cm = confusion_matrix(y_true, y_pred)

# Convert to percentage (row-wise or total)
cm_percent = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] * 100

# Plot
f, ax = plt.subplots(figsize=(4, 3))
sns.heatmap(cm_percent, annot=True, fmt=".1f", cmap="Blues", xticklabels=['Normal', 'OSA'], yticklabels=['Normal', 'OSA'])
ax.set_xlabel("Predicted")
ax.set_ylabel("True")
ax.set_title("Confusion Matrix (% per class)")
plt.show()

f.savefig("confusion_matrix.pdf", dpi=300, bbox_inches='tight')

from sklearn.metrics import roc_curve, roc_auc_score, RocCurveDisplay

fpr, tpr, thresholds = roc_curve(y_true, y_pred_proba)
auc = roc_auc_score(y_true, y_pred_proba)

plt.figure(figsize=(7,4))
RocCurveDisplay(fpr=fpr, tpr=tpr, roc_auc=auc, estimator_name='Model').plot()
plt.title(f"ROC Curve (AUC = {auc:.2f})")
plt.grid(True)
plt.show()

# Step 4: Classification report
from sklearn.metrics import classification_report

print(classification_report(y_true, y_pred, target_names=["No Apnea", "Apnea"]))

for i, (batch_x, batch_y) in enumerate(test_ds):

    plot_examples(batch_x, batch_y, good_ex=True)
    if i>0: break

for i, (batch_x, batch_y) in enumerate(test_ds):

    plot_examples(batch_x, batch_y, good_ex=False, j=i)
    #plot_examples(batch_x, batch_y, good_ex=True, j=i)
    if i>0: break